fix: tie-robust border-ellipse extents so jit gradients are correct - #672

Summary

Fixes a wrong JAX gradient through the border relocator when the lens is centred on the mask's symmetry centre. This is the case for every cell evaluated at the prior medians. Closes #671.

What goes wrong: ellipse_params_via_border_pca_from sizes the relocation ellipse with a = max(|x′|) and b = max(|y′|). A point-symmetric lens (Isothermal + external shear) centred on the symmetry centre of the mask traces the border into exact ± pairs, so both maxima are exact ties. JAX's max gradient shares itself among the entries equal to the result. Under jit, XLA can recompute the operand inside a fused graph with different rounding, so that equality test is unstable. The reverse-mode gradient of b came out 2× or 0.5× its true value, while eager mode was correct. Every pixelized-source likelihood with use_border_relocator=True inherits this.

Measured (jitted Fitness, jax_test imaging, RectangularRTUAdaptImage 39x39, MGE-60 lens light, exact prior medians):

Stage bisection: the traced grid had identical tangents under eager, jit and FD. The relocated grid diverged only under jit reverse mode, and inside it only b (eager 0.5392, jit 1.0908, FD 0.5392).

Fix: a new helper max_at_argmax_from takes the element at the argmax instead of a max reduction, so the whole gradient lands on one point, identically under eager and jit. Values are bit-identical to before. At a genuine tie, any single element's derivative is a valid one-sided derivative of the maximum.

Not changed: the kernel-CDF per-axis min/max and std(axis=0).min() in interpolator/rectangular.py are the same kind of reduction. They showed no eager/jit divergence in axis-aligned and round configurations, which do produce exact ties, so they are left as they are.

About the original #671 symptom (|g| ≈ 2e9 on the HST fiducial): that part is not a library bug. The 141x141 grid at pixelization over-sampling 1 has a point exactly at the Isothermal centre, where the psi softening (a 1e-8 core) makes the true derivative about 1e8, which FD cannot resolve. The evaluation-point fix belongs to the profiling campaign (autolens_profiling#509). Diagnosis is on #671.

API Changes

Bug fix to JAX gradients only; numerical values are unchanged.

• Added autoarray.inversion.mesh.border_relocator.max_at_argmax_from(values, xp=np), a tie-robust maximum.
• ellipse_params_via_border_pca_from uses it for the semi-axes a and b. Same values, corrected jit gradients.
  See full details below.

Test Plan

• New test__max_at_argmax_from__tied_maximum_gradient_lands_on_one_element: one-hot gradient on a tied maximum, eager and jit. jnp.max fails it with a 0.5/0.5 split.
• Full test_autoarray suite: 2268 passed, 4 xfailed.
• End-to-end jitted repro: mass/shear gradients match FD (table above). Repro scripts in Science/scratch/rectangular_gradient_rtu_sparse_2026-11-02/.
• Profiling Witness on the HST fiducial once autolens_profiling#509 moves the fiducial evaluation point off the grid (the odd grid's symmetry centre is a grid point).

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code